import numpy as np
import os

PATCH_COUNT = 400000
COMPONENT_COUNT = 100
n_units = 800
OUT_FOLDER = 'icacodes/'

###
#Load the PCA transformed V1 complex mean subtracted data for the dataset here

#Default is the original ImageNet input patches
pcaTransformed = np.load('imnet/pcaTransformed.npy')[:, :COMPONENT_COUNT]
dataset = 'imnet'

#Mod index: texture naturalistic stimuli
#pcaTransformed = np.load('modindex_nat/pcaTransformed.npy')[:, :COMPONENT_COUNT]
#dataset = 'nat'

#Mod index: texture noise stimuli
#pcaTransformed = np.load('modindex_noise/pcaTransformed.npy')[:, :COMPONENT_COUNT]
#dataset = 'noise'

#Figure-ground stimuli
#pcaTransformed = np.load('fg/pcaTransformed.npy')[:, :COMPONENT_COUNT]
#dataset = 'fg'

#Texture stimuli
#pcaTransformed = np.load('tex/pcaTransformed.npy')[:, :COMPONENT_COUNT]
#dataset = 'tex'

#Line angle stimuli
#pcaTransformed = np.load('line/pcaTransformed.npy')[:, :COMPONENT_COUNT]
#dataset = 'line'

inFile = 'icabasis/basisicacomps' + str(COMPONENT_COUNT) + 'codes' + str(n_units) + 'patches' + str(PATCH_COUNT) + '.npy'
outFile = OUT_FOLDER + dataset + 'icacodesicacomps' + str(COMPONENT_COUNT) + 'codes' + str(n_units) + 'patches' + str(PATCH_COUNT) + '.npy'

if not os.path.exists(OUT_FOLDER):
    os.makedirs(OUT_FOLDER)

icaBasis = np.load(inFile)

#ICA step
codes = np.dot(pcaTransformed, icaBasis.T)

#ReLU step
codes[np.where(codes <= 0.0)] = 0.0

np.save(outFile, codes)
